﻿using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CodeActions;
using Microsoft.CodeAnalysis.CodeFixes;
using Microsoft.CodeAnalysis.CodeRefactorings;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Diagnostics;
using Microsoft.CodeAnalysis.Formatting;
using Microsoft.CodeAnalysis.Simplification;
using Microsoft.CodeAnalysis.Text;
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;

namespace CodeAnalysisRuleCollection.Common.CSharp
{
    /// <example>static bool IsEmpty(string s) { return s.Length == 0; }</example>
    [DiagnosticAnalyzer]
    [ExportDiagnosticAnalyzer(DiagnosticId, LanguageNames.CSharp)]
    internal class CheckParameterForNullDiagnosticAnalyzer : ISyntaxNodeAnalyzer<SyntaxKind>
    {
        #region Description
        public const string DiagnosticId = "CheckParameterForNull";

        public static readonly DiagnosticDescriptor Rule = new DiagnosticDescriptor(
            DiagnosticId,
            "Check parameter for null",
            "Parameter '{0}' is not checked for null before use",
            "Usage",
            DiagnosticSeverity.None
        );

        public virtual ImmutableArray<SyntaxKind> SyntaxKindsOfInterest
        {
            get
            {
                return ImmutableArray.Create(
                    SyntaxKind.MethodDeclaration,
                    // SyntaxKind.ConstructorDeclaration, // TODO: support constructors
                    SyntaxKind.SetAccessorDeclaration,
                    SyntaxKind.AddAccessorDeclaration,
                    SyntaxKind.RemoveAccessorDeclaration
                );
            }
        }

        public ImmutableArray<DiagnosticDescriptor> SupportedDiagnostics { get { return ImmutableArray.Create(Rule); } }
        #endregion

        public virtual void AnalyzeNode(SyntaxNode node, SemanticModel semanticModel, Action<Diagnostic> addDiagnostic, CancellationToken cancellationToken)
        {
            if (
                (var method = node as BaseMethodDeclarationSyntax) != null &&
                method.Body != null &&
                (method.Modifiers.Any(SyntaxKind.PublicKeyword) || method.Modifiers.Any(SyntaxKind.ProtectedKeyword))
            )
            {
                foreach (ParameterSyntax parameterSyntax in method.ParameterList.Parameters)
                {
                    IParameterSymbol symbol = semanticModel.GetDeclaredSymbol(parameterSyntax, cancellationToken);

                    // find member access and invocation expression references
                    if (NeedsCheck(symbol, method.Body, semanticModel, cancellationToken))
                        addDiagnostic(Diagnostic.Create(Rule, parameterSyntax.GetLocation(), symbol.Name));
                }
            }

            if (
                (var method = node as AccessorDeclarationSyntax) != null &&
                method.Body != null
            )
            {
                IParameterSymbol symbol = semanticModel.GetDeclaredSymbol(method, cancellationToken).Parameters.Single();

                // find member access and invocation expression references
                if (NeedsCheck(symbol, method.Body, semanticModel, cancellationToken))
                    addDiagnostic(Diagnostic.Create(Rule, method.Keyword.GetLocation(), symbol.Name));
            }
        }

        private bool NeedsCheck(IParameterSymbol symbol, BlockSyntax block, SemanticModel semanticModel, CancellationToken c)
        {
            // exclude value types
            if (symbol.Type.IsValueType)
                return false;

            // no need to check for initial value, because if it is a reference type, than the only valid constant value is null

            IEnumerable<SyntaxNode> descendants = block.DescendantNodes().ToList();

            // find member access and invocation expression references
            if (
                descendants.OfType<MemberAccessExpressionSyntax>().Any(
                    m => symbol.Equals(semanticModel.GetSymbolInfo(m.Expression, c).Symbol)
                )
                ||
                descendants.OfType<InvocationExpressionSyntax>().Any(
                    m => symbol.Equals(semanticModel.GetSymbolInfo(m.Expression, c).Symbol)
                )
                ||
                descendants.OfType<ElementAccessExpressionSyntax>().Any(
                    m => symbol.Equals(semanticModel.GetSymbolInfo(m.Expression, c).Symbol)
                )
                ||
                descendants.OfType<ForEachStatementSyntax>().Any(
                    m => symbol.Equals(semanticModel.GetSymbolInfo(m.Expression, c).Symbol)
                )
                ||
                descendants.OfType<UsingStatementSyntax>().Any(
                    m => symbol.Equals(semanticModel.GetSymbolInfo(m.Expression, c).Symbol)
                )
                ||
                descendants.OfType<LockStatementSyntax>().Any(
                    m => symbol.Equals(semanticModel.GetSymbolInfo(m.Expression, c).Symbol)
                )
            )
            {
                // find a comparison to the parameter
                if (
                    !descendants.OfType<BinaryExpressionSyntax>().Any(
                        b => (b.IsKind(SyntaxKind.EqualsExpression) || b.IsKind(SyntaxKind.NotEqualsExpression)) &&
                            (symbol.Equals(semanticModel.GetSymbolInfo(b.Left, c).Symbol) || symbol.Equals(semanticModel.GetSymbolInfo(b.Right, c).Symbol))
                    )
                    &&
                    !descendants.OfType<InvocationExpressionSyntax>().Any(
                        b => b.ArgumentList.Arguments.Count > 0 &&
                            b.Expression.ToFullString().Contains("Null") &&
                            symbol.Equals(semanticModel.GetSymbolInfo(b.ArgumentList.Arguments.First().Expression, c).Symbol)
                    )
                )
                    return true;
            }

            return false;
        }
    }

    [ExportCodeFixProvider(CheckParameterForNullDiagnosticAnalyzer.DiagnosticId, LanguageNames.CSharp)]
    internal class CheckParameterForNullCodeFixProvider : ICodeFixProvider
    {
        #region Description
        public IEnumerable<string> GetFixableDiagnosticIds()
        {
            return new[] { CheckParameterForNullDiagnosticAnalyzer.DiagnosticId };
        }
        #endregion

        public async Task<IEnumerable<CodeAction>> GetFixesAsync(Document document, TextSpan span, IEnumerable<Diagnostic> diagnostics, CancellationToken cancellationToken)
        {
            var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);

            var diagnosticSpan = diagnostics.First().Location.SourceSpan;
            var diagnosticNode = root.FindToken(diagnosticSpan.Start).Parent;

            // find the parameter declaration identified by the diagnostic
            if ((var parameterSyntax = diagnosticNode.FirstAncestorOrSelf<ParameterSyntax>()) != null)
            {
                var methodSyntax = parameterSyntax.Parent.Parent as BaseMethodDeclarationSyntax;

                IList<CodeAction> actions = new List<CodeAction>(); // parameter type is always nullable
                actions.Add(CodeAction.Create("Check parameter for null", c => CheckParameterCodeRefactoringProvider.FixAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));
                actions.Add(CodeAction.Create("Check parameter for null (using contract)", c => CheckParameterCodeRefactoringProvider.FixWithContractAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));

                // add additional options by type
                var model = await document.GetSemanticModelAsync(cancellationToken);
                var parameter = model.GetDeclaredSymbol(parameterSyntax, cancellationToken);

                // string
                if (parameter.Type.Equals(model.Compilation.GetSpecialType(SpecialType.System_String)))
                    actions.Add(CodeAction.Create("Check parameter for null or empty string", c => CheckParameterCodeRefactoringProvider.FixWithNullOrEmptyAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));

                return actions;
            }

            // in other cases, when the node is an accessor declaration
            else if ((var accessorSyntax = diagnosticNode.FirstAncestorOrSelf<AccessorDeclarationSyntax>()) != null)
            {
                IList<CodeAction> actions = new List<CodeAction>(); // parameter type is always nullable
                actions.Add(CodeAction.Create("Check parameter for null", c => CheckParameterCodeRefactoringProvider.FixAsync(root, document, "value", accessorSyntax.Body, c)));

                // add additional options by type
                var model = await document.GetSemanticModelAsync(cancellationToken);
                var parameter = model.GetDeclaredSymbol(parameterSyntax, cancellationToken);

                // string
                if (parameter.Type.Equals(model.Compilation.GetSpecialType(SpecialType.System_String)))
                    actions.Add(CodeAction.Create("Check parameter for null or empty string", c => CheckParameterCodeRefactoringProvider.FixWithNullOrEmptyAsync(root, document, "value", accessorSyntax.Body, c)));

                return actions;
            }

            return null;
        }
    }

    [ExportCodeRefactoringProvider("CheckParameter", LanguageNames.CSharp)]
    internal class CheckParameterCodeRefactoringProvider : ICodeRefactoringProvider
    {
        public async Task<IEnumerable<CodeAction>> GetRefactoringsAsync(Document document, TextSpan span, CancellationToken cancellationToken)
        {
            var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);

            var diagnosticNode = root.FindToken(span.Start).Parent;

            // find the parameter declaration identified by the diagnostic
            if ((var parameterSyntax = diagnosticNode.FirstAncestorOrSelf<ParameterSyntax>()) != null)
            {
                var methodSyntax = parameterSyntax.Parent.Parent as BaseMethodDeclarationSyntax;

                IList<CodeAction> actions = new List<CodeAction>();

                // provide additional options by type
                var model = await document.GetSemanticModelAsync(cancellationToken);
                var parameter = model.GetDeclaredSymbol(parameterSyntax, cancellationToken);

                if (!parameter.Type.IsValueType)
                {
                    actions.Add(CodeAction.Create("Check parameter for null", c => FixAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));
                    actions.Add(CodeAction.Create("Check parameter for null (using contract)", c => FixWithContractAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));
                }

                // string
                if (parameter.Type.Equals(model.Compilation.GetSpecialType(SpecialType.System_String)))
                    actions.Add(CodeAction.Create("Check parameter for null or empty string", c => FixWithNullOrEmptyAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));

                // int
                else if (parameter.Type.Equals(model.Compilation.GetSpecialType(SpecialType.System_Int32)))
                    actions.Add(CodeAction.Create("Check parameter for negative", c => FixWithNotNegativeAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));

                // DateTime
                else if (parameter.Type.Equals(model.Compilation.GetSpecialType(SpecialType.System_DateTime)))
                    actions.Add(CodeAction.Create("Check parameter for later than today", c => FixWithOlderThanTodayAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));

                // IntPtr
                else if (parameter.Type.Equals(model.Compilation.GetSpecialType(SpecialType.System_IntPtr)))
                    actions.Add(CodeAction.Create("Check parameter for IntPtr.Zero", c => FixWithNotIntPtrZeroAsync(root, document, parameterSyntax.Identifier.ValueText, methodSyntax.Body, c)));
                
                return actions;
            }

            return null;
        }

        internal static async Task<Document> FixAsync(SyntaxNode root, Document document, string parameterName, BlockSyntax blockSyntax, CancellationToken cancellationToken)
        {
            // TODO: Find other checks and insert into the right place

            var newRoot = root.ReplaceNode(
                blockSyntax,
                blockSyntax.WithStatements(
                    blockSyntax.Statements.Insert(0,
                        SyntaxFactory.IfStatement(
                            SyntaxFactory.BinaryExpression(SyntaxKind.EqualsExpression,
                                SyntaxFactory.IdentifierName(parameterName),
                                SyntaxFactory.LiteralExpression(SyntaxKind.NullLiteralExpression)
                            ),
                            SyntaxFactory.ThrowStatement(
                                SyntaxFactory.ObjectCreationExpression(
                                    SyntaxFactory.ParseTypeName("System.ArgumentNullException")
                                        .WithAdditionalAnnotations(Simplifier.Annotation),
                                    SyntaxFactory.ArgumentList(
                                        SyntaxFactory.SeparatedList(
                                            new ArgumentSyntax[]
                                            {
                                                SyntaxFactory.Argument(
                                                    SyntaxFactory.LiteralExpression(SyntaxKind.StringLiteralExpression, SyntaxFactory.Literal(parameterName))
                                                )
                                            }
                                        )
                                    ),
                                    null
                                )
                            )
                        )
                            .WithTrailingTrivia(
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine),
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine)
                            )
                            .WithAdditionalAnnotations(Formatter.Annotation)
                    )
                )
            );

            return document.WithSyntaxRoot(newRoot);
        }

        internal static async Task<Document> FixWithContractAsync(SyntaxNode root, Document document, string parameterName, BlockSyntax blockSyntax, CancellationToken cancellationToken)
        {
            // TODO: Find other checks and insert into the right place

            var newRoot = root.ReplaceNode(
                blockSyntax,
                blockSyntax.WithStatements(
                    blockSyntax.Statements.Insert(0,
                        SyntaxFactory.ExpressionStatement(
                            SyntaxFactory.InvocationExpression(
                                SyntaxFactory.ParseName("Contract.Requires")
                                    .WithAdditionalAnnotations(Simplifier.Annotation),
                                SyntaxFactory.ArgumentList(
                                    SyntaxFactory.SeparatedList(
                                        new ArgumentSyntax[]
                                        {
                                            SyntaxFactory.Argument(
                                                SyntaxFactory.BinaryExpression(SyntaxKind.NotEqualsExpression,
                                                    SyntaxFactory.IdentifierName(parameterName),
                                                    SyntaxFactory.LiteralExpression(SyntaxKind.NullLiteralExpression)
                                                )
                                            )
                                        }
                                    )
                                )
                            )
                                .WithAdditionalAnnotations(Formatter.Annotation)
                        )
                            .WithTrailingTrivia(
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine),
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine)
                            )
                    )
                )
            );

            // add using to Code Contracts
            newRoot = (newRoot as CompilationUnitSyntax).WithUsing("System.Diagnostics.Contracts");

            return document.WithSyntaxRoot(newRoot);
        }

        internal static async Task<Document> FixWithNullOrEmptyAsync(SyntaxNode root, Document document, string parameterName, BlockSyntax blockSyntax, CancellationToken cancellationToken)
        {
            // TODO: Find other checks and insert into the right place

            var newRoot = root.ReplaceNode(
                blockSyntax,
                blockSyntax.WithStatements(
                    blockSyntax.Statements.Insert(0,
                        SyntaxFactory.IfStatement(
                            SyntaxFactory.InvocationExpression(
                                SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
                                    SyntaxFactory.ParseName("System.String")
                                        .WithAdditionalAnnotations(Simplifier.Annotation),
                                    SyntaxFactory.IdentifierName("IsNullOrEmpty")
                                ),
                                SyntaxFactory.ArgumentList(
                                    SyntaxFactory.SeparatedList(
                                        new ArgumentSyntax[]
                                        {
                                            SyntaxFactory.Argument(
                                                SyntaxFactory.IdentifierName(parameterName)
                                            )
                                        }
                                    )
                                )
                            ),
                            SyntaxFactory.ThrowStatement(
                                SyntaxFactory.ObjectCreationExpression(
                                    SyntaxFactory.ParseTypeName("System.ArgumentNullException")
                                        .WithAdditionalAnnotations(Simplifier.Annotation),
                                    SyntaxFactory.ArgumentList(
                                        SyntaxFactory.SeparatedList(
                                            new ArgumentSyntax[]
                                            {
                                                SyntaxFactory.Argument(
                                                    SyntaxFactory.LiteralExpression(SyntaxKind.StringLiteralExpression, SyntaxFactory.Literal(parameterName))
                                                )
                                            }
                                        )
                                    ),
                                    null
                                )
                            )
                        )
                            .WithTrailingTrivia(
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine),
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine)
                            )
                            .WithAdditionalAnnotations(Formatter.Annotation)
                    )
                )
            );

            return document.WithSyntaxRoot(newRoot);
        }

        internal static async Task<Document> FixWithNotNegativeAsync(SyntaxNode root, Document document, string parameterName, BlockSyntax blockSyntax, CancellationToken cancellationToken)
        {
            // TODO: Find other checks and insert into the right place

            var newRoot = root.ReplaceNode(
                blockSyntax,
                blockSyntax.WithStatements(
                    blockSyntax.Statements.Insert(0,
                        SyntaxFactory.IfStatement(
                            SyntaxFactory.BinaryExpression(SyntaxKind.LessThanExpression,
                                SyntaxFactory.IdentifierName(parameterName),
                                SyntaxFactory.LiteralExpression(SyntaxKind.NumericLiteralExpression, SyntaxFactory.Literal(0))
                            ),
                            SyntaxFactory.ThrowStatement(
                                SyntaxFactory.ObjectCreationExpression(
                                    SyntaxFactory.ParseTypeName("System.ArgumentOutOfRangeException")
                                        .WithAdditionalAnnotations(Simplifier.Annotation),
                                    SyntaxFactory.ArgumentList(
                                        SyntaxFactory.SeparatedList(
                                            new ArgumentSyntax[]
                                            {
                                                SyntaxFactory.Argument(
                                                    SyntaxFactory.LiteralExpression(SyntaxKind.StringLiteralExpression, SyntaxFactory.Literal(parameterName))
                                                )
                                            }
                                        )
                                    ),
                                    null
                                )
                            )
                        )
                            .WithTrailingTrivia(
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine),
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine)
                            )
                            .WithAdditionalAnnotations(Formatter.Annotation)
                    )
                )
            );

            return document.WithSyntaxRoot(newRoot);
        }

        internal static async Task<Document> FixWithOlderThanTodayAsync(SyntaxNode root, Document document, string parameterName, BlockSyntax blockSyntax, CancellationToken cancellationToken)
        {
            // TODO: Find other checks and insert into the right place

            var newRoot = root.ReplaceNode(
                blockSyntax,
                blockSyntax.WithStatements(
                    blockSyntax.Statements.Insert(0,
                        SyntaxFactory.IfStatement(
                            SyntaxFactory.BinaryExpression(SyntaxKind.GreaterThanExpression,
                                SyntaxFactory.IdentifierName(parameterName),
                                SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
                                    SyntaxFactory.ParseTypeName("System.DateTime")
                                        .WithAdditionalAnnotations(Simplifier.Annotation),
                                    SyntaxFactory.IdentifierName("Now")
                                )
                            ),
                            SyntaxFactory.ThrowStatement(
                                SyntaxFactory.ObjectCreationExpression(
                                    SyntaxFactory.ParseTypeName("System.ArgumentOutOfRangeException")
                                        .WithAdditionalAnnotations(Simplifier.Annotation),
                                    SyntaxFactory.ArgumentList(
                                        SyntaxFactory.SeparatedList(
                                            new ArgumentSyntax[]
                                            {
                                                SyntaxFactory.Argument(
                                                    SyntaxFactory.LiteralExpression(SyntaxKind.StringLiteralExpression, SyntaxFactory.Literal(parameterName))
                                                )
                                            }
                                        )
                                    ),
                                    null
                                )
                            )
                        )
                            .WithTrailingTrivia(
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine),
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine)
                            )
                            .WithAdditionalAnnotations(Formatter.Annotation)
                    )
                )
            );

            return document.WithSyntaxRoot(newRoot);
        }

        internal static async Task<Document> FixWithNotIntPtrZeroAsync(SyntaxNode root, Document document, string parameterName, BlockSyntax blockSyntax, CancellationToken cancellationToken)
        {
            // TODO: Find other checks and insert into the right place
            
            var newRoot = root.ReplaceNode(
                blockSyntax,
                blockSyntax.WithStatements(
                    blockSyntax.Statements.Insert(0,
                        SyntaxFactory.IfStatement(
                            SyntaxFactory.BinaryExpression(SyntaxKind.EqualsExpression,
                                SyntaxFactory.IdentifierName(parameterName),
                                SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
                                    SyntaxFactory.ParseTypeName("System.IntPtr")
                                        .WithAdditionalAnnotations(Simplifier.Annotation),
                                    SyntaxFactory.IdentifierName("Zero")
                                )
                            ),
                            SyntaxFactory.ThrowStatement(
                                SyntaxFactory.ObjectCreationExpression(
                                    SyntaxFactory.ParseTypeName("System.ArgumentNullException")
                                        .WithAdditionalAnnotations(Simplifier.Annotation),
                                    SyntaxFactory.ArgumentList(
                                        SyntaxFactory.SeparatedList(
                                            new ArgumentSyntax[]
                                            {
                                                SyntaxFactory.Argument(
                                                    SyntaxFactory.LiteralExpression(SyntaxKind.StringLiteralExpression, SyntaxFactory.Literal(parameterName))
                                                )
                                            }
                                        )
                                    ),
                                    null
                                )
                            )
                        )
                            .WithTrailingTrivia(
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine),
                                SyntaxFactory.SyntaxTrivia(SyntaxKind.EndOfLineTrivia, Environment.NewLine)
                            )
                            .WithAdditionalAnnotations(Formatter.Annotation)
                    )
                )
            );

            return document.WithSyntaxRoot(newRoot);
        }
    }
}